package main.naiveBayes;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;

import edu.stanford.smi.protegex.owl.model.OWLProperty;

import y.base.Edge;
import y.base.Node;

import main.graph.FormTreeWithMapping;
import main.graph.ValuedGraph;
import main.model.Attribute;
import main.model.Cardinality;
import main.model.DummyEdge;
import main.model.EdgeGivenPath;
import main.model.EdgePriors;
import main.model.OntoEle;
import main.model.OntoNodePair;
import main.model.OntoPath;
import main.model.Relationship;
import main.model.Tree2OntoMapping;

/**
 * Learn the parameters for the naive Bayesian method:
 * 1. priors. 2. likelihoods
 * 
 * @author YuanAn
 *
 */
public class LearnParameter {

	/**
	 * learn the priors of each ontology edge from the training data, a set of form tree with mappings.
	 * 
	 * (nc+1)/(n+m): nc is the number of times the edge appeared in the training data; n is the number of 
	 * total edges appeared in the training data; m is the number of edges in the ontology
	 * 
	 * @param trees
	 */
	public static HashMap<OntoEle, Double> learnPrior(ArrayList<FormTreeWithMapping> trees, ValuedGraph ontoGraph){
		
		int nodeCount = ontoGraph.N();
		int numPairs = (nodeCount * (nodeCount - 1)) / 2;
		
		int numEdges = ontoGraph.E();
		
		System.out.println("The number of pairs of nodes in the ontology: " + numPairs);
		System.out.println("The number of edges in the ontology: " + numEdges);
		
		HashMap<OntoEle, Integer> edgeCount = new HashMap<OntoEle, Integer>();
		
		int totalEdgeCount = 0;
		
		for(FormTreeWithMapping tree: trees){
			Tree2OntoMapping mapping = tree.getMapping();
			
			HashMap<Edge, OntoPath> edge2onto = mapping.getEdge2Onto();
			
			Collection<OntoPath> paths = edge2onto.values();
			
			Iterator<OntoPath> pit = paths.iterator();
			
			//for each ontoEdge, count the number of appearances in the training data
			while(pit.hasNext()){
				OntoPath path = pit.next();
				ArrayList<OntoEle> pathEles = path.getPath();
				for(OntoEle ele: pathEles){
					if(ele instanceof Attribute || ele instanceof Relationship){
						
						//total number of the edges in the training data
						totalEdgeCount ++;
						//System.out.println("LearnParameter: learnPriorAndLikelihood: OntoEle's name: " + ele.display());
						Integer countInteger = edgeCount.get(ele);
						if(countInteger == null){
							edgeCount.put(ele, 1);
						}
						else{
							int count = countInteger.intValue() + 1;
							edgeCount.put(ele, count);
						}
					}
				}
			}
		}
		
		//hold the prior for each onto edge in the training data
		HashMap<OntoEle, Double> edgePriors = new HashMap<OntoEle, Double>();
		
		Iterator<OntoEle> ontoEleIt = edgeCount.keySet().iterator();
		while(ontoEleIt.hasNext()){
			OntoEle ele = ontoEleIt.next();
			int count =  edgeCount.get(ele).intValue();
			
			//compute the logarithm of the prior: (nc+1)/(n+m)
			double prior = Math.log(((double)(count + 1)) / (totalEdgeCount + numEdges));
			
			edgePriors.put(ele, prior);
		}
		
		//add the prior for all the other edges that are not in the training set 1/(n+m)
		double dummyPrior = Math.log(((double)(1)) / (totalEdgeCount + numEdges));
		edgePriors.put(new DummyEdge(), dummyPrior);
		
		System.out.println("The total number of edges in the training data: " + totalEdgeCount);
		
		return edgePriors;
	}
	
	/**
	 * learn the likelihood of an edge given a path
	 * 
	 * (nc+1)/(n+m): nc is the number of times the edge appeared when the path p appeared; n is the number of 
	 * the times the path p appeared; m is the number of pairs of nodes in the ontology
	 * 
	 * @param trees
	 * @param ontoGraph
	 * @return
	 */
	public static ArrayList<ArrayList<EdgeGivenPath>> learnLikelihood(ArrayList<FormTreeWithMapping> trees, ValuedGraph ontoGraph){
		
		ArrayList<ArrayList<EdgeGivenPath>>  ans = new ArrayList<ArrayList<EdgeGivenPath>>();
		
		int nodeCount = ontoGraph.N();
		int numPairs = (nodeCount * (nodeCount - 1)) / 2;
		
		//System.out.println("The number of pairs of nodes: " + numPairs);
		
		//this information may not be used.
		HashMap<OntoPath, Integer> pathCount = new HashMap<OntoPath, Integer>();
		
		HashMap<OntoPath, ArrayList<EdgeGivenPath>> pathToEdges = new HashMap<OntoPath, ArrayList<EdgeGivenPath>>();
		
		
		for(FormTreeWithMapping tree: trees){
			Tree2OntoMapping mapping = tree.getMapping();
			
			HashMap<Edge, OntoPath> edge2onto = mapping.getEdge2Onto();
			
			Collection<OntoPath> paths = edge2onto.values();
			
			Iterator<OntoPath> pit = paths.iterator();
			
			//hold the paths in the tree
			ArrayList<OntoPath> pathsInTree = new ArrayList<OntoPath>();
			//hold the pairs of nodes in the tree
			ArrayList<OntoNodePair> pairsInTree = new ArrayList<OntoNodePair>();
			
			while(pit.hasNext()){
				OntoPath path = pit.next();
				
				pathsInTree.add(path);
			
				//check whether the path has been counted.
				boolean foundPath = false;
				Iterator<OntoPath> existingPathIt = pathCount.keySet().iterator();
				while(existingPathIt.hasNext()){
					OntoPath existingPath = existingPathIt.next();
					if(existingPath.isSamePath(path)){
						//if the path has been counted
						int count = pathCount.get(existingPath).intValue() + 1;
						pathCount.put(existingPath, count);
						foundPath = true;
						break;
					}
				}
				if(!foundPath){//the path hasn't been counted
					pathCount.put(path, 1);
				}
				
				//count the occurrence of the pair of nodes associated with this path
				OntoEle first = path.getFirst();
				OntoEle second = path.getLast();
				OntoNodePair newPair = new OntoNodePair(first, second);
				
				Node firstNode = path.getFirstNode();
				Node lastNode = path.getLastNode();
				
				newPair.setFirstNode(firstNode);
				newPair.setSecondNode(lastNode);
				
				pairsInTree.add(newPair);
			}
			
			//now, for each path and a pair of nodes, count...
			for(OntoPath aPath: pathsInTree){
				
				ArrayList<EdgeGivenPath> edges = new ArrayList<EdgeGivenPath>();
				
				boolean foundPath = false;
			 	Iterator<OntoPath> existingPathIt = pathToEdges.keySet().iterator();
				while(existingPathIt.hasNext()){
					OntoPath existingPath = existingPathIt.next();
					if(existingPath.isSamePath(aPath)){
						//if the path has been associated with edges
						edges = pathToEdges.get(existingPath);
						foundPath = true;
						break;
					}
				}
				if(!foundPath){
					
					//add a dummy edge representing all the other edges
					OntoNodePair dummyPair = new OntoNodePair();
					EdgeGivenPath dummy = new EdgeGivenPath(dummyPair, aPath);
					edges.add(dummy);
					
					pathToEdges.put(aPath, edges);
				}
				
				for(OntoNodePair aPair: pairsInTree){
							
					//check whether the pair of nodes exist
					boolean foundPair = false;
					for(EdgeGivenPath edge: edges){
						if(edge.getEdge().isSame(aPair)){
							edge.incrementEdgeCount();
							foundPair = true;
							break;
						}
					}
					if(!foundPair){
						EdgeGivenPath newEdge = new EdgeGivenPath(aPair, aPath);
						newEdge.incrementEdgeCount();
						edges.add(newEdge);
					}
				}
			}
		}
		
	
		//test
		Iterator<OntoPath> pathCountIt = pathCount.keySet().iterator();
		while(pathCountIt.hasNext()){
			OntoPath path = pathCountIt.next();
			Integer count = pathCount.get(path);
			System.out.println(path.display() + " :=+ " + count);
		}
		
		ans.addAll(pathToEdges.values());
		
		int totalPathEdgeCount = 0; //the total number of pairs of edge and path in the training data
		for(ArrayList<EdgeGivenPath> edges: ans){
			int totalPathCount = 0;
			//count the total number of paths for all edges
			for(EdgeGivenPath edge: edges)
				totalPathCount += edge.getEdgeCount();
			//compute conditional prob
			for(EdgeGivenPath edge: edges){
				edge.setPathCount(totalPathCount);
				
				//compute the logarithm of the conditionalProb: (nc+1)/(n+m)
				double conditionalProb = Math.log(((double)(edge.getEdgeCount() + 1)) / (edge.getPathCount() + numPairs));
				edge.setConditionalProb(conditionalProb);
			}
			
			totalPathEdgeCount += totalPathCount;
		}
		
		//add a dummy EdgeGivenPath for all other paths not in training data
		EdgeGivenPath dummyEdgeGivenPath = new EdgeGivenPath();
		dummyEdgeGivenPath.setPathCount(totalPathEdgeCount);
		//compute the logarithm of the conditionalProb: (nc+1)/(n+m)
		double conditionalProb = Math.log(((double)(dummyEdgeGivenPath.getEdgeCount() + 1)) / (dummyEdgeGivenPath.getPathCount() + numPairs));
		dummyEdgeGivenPath.setConditionalProb(conditionalProb);
		ArrayList<EdgeGivenPath> dummyArray = new ArrayList<EdgeGivenPath>();
		dummyArray.add(dummyEdgeGivenPath);
		ans.add(dummyArray);
		
		return ans;
	}
	
	/**
	 * learn the priors of each ontology edge from the training data, a set of form tree with mappings.
	 * 
	 * (nc+1)/(n+m): nc is the number of times the edge appeared in the training data; n is the number of 
	 * total edges appeared in the training data; m is the number of edges in the ontology
	 * 
	 * @param trees
	 */
	public static EdgePriors learnEdgePrior(ArrayList<FormTreeWithMapping> trees, ValuedGraph ontoGraph){
		
		int nodeCount = ontoGraph.N();
		int numPairs = (nodeCount * (nodeCount - 1)) / 2;
		
		int numEdges = ontoGraph.E();
		
		System.out.println("The number of pairs of nodes in the ontology: " + numPairs);
		System.out.println("The number of edges in the ontology: " + numEdges);
		
		HashMap<Edge, Integer> edgeCount = new HashMap<Edge, Integer>();
		
		int totalEdgeCount = 0;
		
		for(FormTreeWithMapping tree: trees){
			Tree2OntoMapping mapping = tree.getMapping();
			
			HashMap<Edge, OntoPath> edge2onto = mapping.getEdge2Onto();
			
			Collection<OntoPath> paths = edge2onto.values();
			
			Iterator<OntoPath> pit = paths.iterator();
			
			//for each ontoEdge, count the number of appearances in the training data
			while(pit.hasNext()){
				OntoPath path = pit.next();
				ArrayList<Edge> pathEdges = path.getEdgePath();
				for(Edge e: pathEdges){
						
					//total number of the edges in the training data
					totalEdgeCount ++;
					//System.out.println("LearnParameter: learnPriorAndLikelihood: OntoEle's name: " + ele.display());
					Integer countInteger = edgeCount.get(e);
					if(countInteger == null){
						edgeCount.put(e, 1);
					}
					else{
						int count = countInteger.intValue() + 1;
						edgeCount.put(e, count);
					}
				}
			}
		}
		
		//hold the prior for each onto edge in the training data
		EdgePriors edgePriors = new EdgePriors();
		
		Iterator<Edge> ontoEdgeIt = edgeCount.keySet().iterator();
		while(ontoEdgeIt.hasNext()){
			Edge e = ontoEdgeIt.next();
			int count =  edgeCount.get(e).intValue();
			
			//compute the logarithm of the prior: (nc+1)/(n+m)
			double prior = Math.log(((double)(count + 1)) / (totalEdgeCount + numEdges));
			
			edgePriors.put(e, prior);
		}
		
		//add the prior for all the other edges that are not in the training set 1/(n+m)
		double dummyPrior = Math.log(((double)(1)) / (totalEdgeCount + numEdges));
		edgePriors.setDummyPrior(dummyPrior);
		
		System.out.println("The total number of edges in the training data: " + totalEdgeCount);
		
		return edgePriors;
	}
	
}
